package org.example.common.azure;

import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.PropertyNamingStrategy;
import com.theokanning.openai.DeleteResult;
import com.theokanning.openai.OpenAiError;
import com.theokanning.openai.OpenAiHttpException;
import com.theokanning.openai.completion.CompletionChunk;
import com.theokanning.openai.completion.CompletionRequest;
import com.theokanning.openai.completion.CompletionResult;
import com.theokanning.openai.completion.chat.ChatCompletionChunk;
import com.theokanning.openai.completion.chat.ChatCompletionRequest;
import com.theokanning.openai.completion.chat.ChatCompletionResult;
import com.theokanning.openai.edit.EditRequest;
import com.theokanning.openai.edit.EditResult;
import com.theokanning.openai.embedding.EmbeddingRequest;
import com.theokanning.openai.embedding.EmbeddingResult;
import com.theokanning.openai.file.File;
import com.theokanning.openai.finetune.FineTuneEvent;
import com.theokanning.openai.finetune.FineTuneRequest;
import com.theokanning.openai.finetune.FineTuneResult;
import com.theokanning.openai.image.CreateImageEditRequest;
import com.theokanning.openai.image.CreateImageRequest;
import com.theokanning.openai.image.CreateImageVariationRequest;
import com.theokanning.openai.image.ImageResult;
import com.theokanning.openai.model.Model;
import com.theokanning.openai.moderation.ModerationRequest;
import com.theokanning.openai.moderation.ModerationResult;
import io.reactivex.BackpressureStrategy;
import io.reactivex.Flowable;
import io.reactivex.Single;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.example.common.azure.pojo.AzureCreateImageRequest;
import org.example.common.azure.pojo.AzureCreateImageResponseNotStarted;
import org.example.common.azure.pojo.AzureCreateImageResult;
import org.example.common.azure.pojo.AzureCreateImageResultNew;
import org.example.common.utils.CP_SleepUtils;
import org.example.common.utils.OkHttpUtil;
import retrofit2.Call;
import retrofit2.HttpException;
import retrofit2.Response;
import retrofit2.Retrofit;
import retrofit2.adapter.rxjava2.RxJava2CallAdapterFactory;
import retrofit2.converter.jackson.JacksonConverterFactory;

import java.io.IOException;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * <h1></h1>
 * Created by hanqf on 2023/4/24 13:54.
 */

@Slf4j
public class AzureOpenAiService {

    private static final String BASE_URL = "https://yyyyy.openai.azure.com/";
    private static final String API_KEY = "xxxxx";

    /**
     * 聊天模块 model-id
     */
    private static final String CHAT_MODEL_ID = "lexing-gpt-35-turbo";
    /**
     * 聊天模块 api-version
     * azure最新api版本接口文件：
     * https://github.com/Azure/azure-rest-api-specs/blob/main/specification/cognitiveservices/data-plane/AzureOpenAI/inference/preview/2023-06-01-preview/inference.json
     */
//    private static final String CHAT_API_VERSION = "2023-03-15-preview";
    private static final String CHAT_API_VERSION = "2023-06-01-preview";

    /**
     * 创建图片模块 api-version
     */
    private static final String CREATE_IMAGE_API_VERSION = "2022-08-03-preview";
    private static final String CREATE_IMAGE_API_VERSION_NEW = "2023-06-01-preview";


    private static final ObjectMapper mapper = defaultObjectMapper();
    private static final AzureOpenAiApi api = api();
    private static final ExecutorService executorService = defaultClient().dispatcher().executorService();

    public static ObjectMapper defaultObjectMapper() {
        ObjectMapper mapper = new ObjectMapper();
        mapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
        mapper.setSerializationInclusion(JsonInclude.Include.NON_NULL);
        mapper.setPropertyNamingStrategy(PropertyNamingStrategy.SNAKE_CASE);
        return mapper;
    }

    public static OkHttpClient defaultClient() {
        return OkHttpUtil.client.newBuilder()
                .addInterceptor(new AzureAuthenticationInterceptor(API_KEY))
                .connectionPool(new ConnectionPool(5, 1, TimeUnit.SECONDS))
                .build();
        //return new OkHttpClient.Builder()
        //        .addInterceptor(new AzureAuthenticationInterceptor(API_KEY))
        //        .connectionPool(new ConnectionPool(5, 1, TimeUnit.SECONDS))
        //        .readTimeout(5, TimeUnit.MINUTES)
        //        .build();
    }


    public static Retrofit defaultRetrofit() {
        return new Retrofit.Builder()
                .baseUrl(BASE_URL)
                .client(defaultClient())
                .addConverterFactory(JacksonConverterFactory.create(defaultObjectMapper()))
                .addCallAdapterFactory(RxJava2CallAdapterFactory.create())
                .build();
    }


    public static AzureOpenAiApi api() {
        return defaultRetrofit().create(AzureOpenAiApi.class);
    }


    public List<Model> listModels() {
        return execute(api.listModels()).data;
    }

    public Model getModel(String modelId) {
        return execute(api.getModel(modelId));
    }

    public CompletionResult createCompletion(CompletionRequest request) {
        return execute(api.createCompletion(request, CHAT_MODEL_ID, CHAT_API_VERSION));
    }

    public Flowable<CompletionChunk> streamCompletion(CompletionRequest request) {
        request.setStream(true);

        return stream(api.createCompletionStream(request, CHAT_MODEL_ID, CHAT_API_VERSION), CompletionChunk.class);
    }

    public ChatCompletionResult createChatCompletion(ChatCompletionRequest request) {
        return execute(api.createChatCompletion(request, CHAT_MODEL_ID, CHAT_API_VERSION));
    }

    public Flowable<ChatCompletionChunk> streamChatCompletion(ChatCompletionRequest request) {
        request.setStream(true);

        return stream(api.createChatCompletionStream(request, CHAT_MODEL_ID, CHAT_API_VERSION), ChatCompletionChunk.class);
    }

    public AzureCreateImageResult azureCreateImageResult(AzureCreateImageRequest request) {
        final AzureCreateImageResponseNotStarted notStarted = execute(api.createAuzreImageNotStarted(request, CREATE_IMAGE_API_VERSION));
        AzureCreateImageResult imageResult;
        CP_SleepUtils.sleep(6, TimeUnit.SECONDS);
        imageResult = execute(api.createAuzreImageByNotStartedId(notStarted.getId(), CREATE_IMAGE_API_VERSION));
        int waitTime = 1; //暂时控制等待时间为2分钟
        while (!"Succeeded".equals(imageResult.getStatus()) && waitTime < 20) {
            CP_SleepUtils.sleep(6, TimeUnit.SECONDS);
            waitTime++;
            imageResult = execute(api.createAuzreImageByNotStartedId(notStarted.getId(), CREATE_IMAGE_API_VERSION));
        }
        return imageResult;
    }

    public AzureCreateImageResult azureCreateImageResult2(AzureCreateImageRequest request) {
        AzureCreateImageResult imageResult = null;
        final Response<AzureCreateImageResponseNotStarted> response = api.createAuzreImageNotStartedResponse(request, CREATE_IMAGE_API_VERSION);
        if (response.isSuccessful()) {
            AzureCreateImageResponseNotStarted notStarted = response.body();
            final String retry_after = response.headers().get("Retry-after") == null ? "6" : response.headers().get("Retry-after");
            CP_SleepUtils.sleep(Long.parseLong(retry_after), TimeUnit.SECONDS);
            imageResult = execute(api.createAuzreImageByNotStartedId(notStarted.getId(), CREATE_IMAGE_API_VERSION));
            int waitTime = 1; //暂时控制等待时间为2分钟
            while (!"Succeeded".equals(imageResult.getStatus()) && waitTime < 20) {
                CP_SleepUtils.sleep(Long.parseLong(retry_after), TimeUnit.SECONDS);
                waitTime++;
                imageResult = execute(api.createAuzreImageByNotStartedId(notStarted.getId(), CREATE_IMAGE_API_VERSION));
            }
        }
        return imageResult;
    }


    /**
     * 创建图像新接口
    */
    public AzureCreateImageResultNew azureCreateImageResultNew(CreateImageRequest request) {
        final AzureCreateImageResponseNotStarted notStarted = execute(api.createAuzreImageNotStartedNew(request, CREATE_IMAGE_API_VERSION_NEW));
        AzureCreateImageResultNew imageResult;
        CP_SleepUtils.sleep(6, TimeUnit.SECONDS);
        imageResult = execute(api.createAuzreImageByNotStartedIdNew(notStarted.getId(), CREATE_IMAGE_API_VERSION_NEW));
        int waitTime = 1; //暂时控制等待时间为2分钟
        while (!"succeeded".equals(imageResult.getStatus()) && waitTime < 20) {
            CP_SleepUtils.sleep(6, TimeUnit.SECONDS);
            waitTime++;
            imageResult = execute(api.createAuzreImageByNotStartedIdNew(notStarted.getId(), CREATE_IMAGE_API_VERSION_NEW));
        }
        return imageResult;
    }

    public EditResult createEdit(EditRequest request) {
        return execute(api.createEdit(request));
    }

    public EmbeddingResult createEmbeddings(EmbeddingRequest request) {
        return execute(api.createEmbeddings(request));
    }

    public List<File> listFiles() {
        return execute(api.listFiles()).data;
    }

    public File uploadFile(String purpose, String filepath) {
        java.io.File file = new java.io.File(filepath);
        RequestBody purposeBody = RequestBody.create(okhttp3.MultipartBody.FORM, purpose);
        RequestBody fileBody = RequestBody.create(MediaType.parse("text"), file);
        MultipartBody.Part body = MultipartBody.Part.createFormData("file", filepath, fileBody);

        return execute(api.uploadFile(purposeBody, body));
    }

    public DeleteResult deleteFile(String fileId) {
        return execute(api.deleteFile(fileId));
    }

    public File retrieveFile(String fileId) {
        return execute(api.retrieveFile(fileId));
    }

    public FineTuneResult createFineTune(FineTuneRequest request) {
        return execute(api.createFineTune(request));
    }

    public CompletionResult createFineTuneCompletion(CompletionRequest request) {
        return execute(api.createFineTuneCompletion(request));
    }

    public List<FineTuneResult> listFineTunes() {
        return execute(api.listFineTunes()).data;
    }

    public FineTuneResult retrieveFineTune(String fineTuneId) {
        return execute(api.retrieveFineTune(fineTuneId));
    }

    public FineTuneResult cancelFineTune(String fineTuneId) {
        return execute(api.cancelFineTune(fineTuneId));
    }

    public List<FineTuneEvent> listFineTuneEvents(String fineTuneId) {
        return execute(api.listFineTuneEvents(fineTuneId)).data;
    }

    public DeleteResult deleteFineTune(String fineTuneId) {
        return execute(api.deleteFineTune(fineTuneId));
    }

    public ImageResult createImage(CreateImageRequest request) {
        return execute(api.createImage(request));
    }

    public ImageResult createImageEdit(CreateImageEditRequest request, String imagePath, String maskPath) {
        java.io.File image = new java.io.File(imagePath);
        java.io.File mask = null;
        if (maskPath != null) {
            mask = new java.io.File(maskPath);
        }
        return createImageEdit(request, image, mask);
    }

    public ImageResult createImageEdit(CreateImageEditRequest request, java.io.File image, java.io.File mask) {
        RequestBody imageBody = RequestBody.create(MediaType.parse("image"), image);

        MultipartBody.Builder builder = new MultipartBody.Builder()
                .setType(MediaType.get("multipart/form-data"))
                .addFormDataPart("prompt", request.getPrompt())
                .addFormDataPart("size", request.getSize())
                .addFormDataPart("response_format", request.getResponseFormat())
                .addFormDataPart("image", "image", imageBody);

        if (request.getN() != null) {
            builder.addFormDataPart("n", request.getN().toString());
        }

        if (mask != null) {
            RequestBody maskBody = RequestBody.create(MediaType.parse("image"), mask);
            builder.addFormDataPart("mask", "mask", maskBody);
        }

        return execute(api.createImageEdit(builder.build()));
    }

    public ImageResult createImageVariation(CreateImageVariationRequest request, String imagePath) {
        java.io.File image = new java.io.File(imagePath);
        return createImageVariation(request, image);
    }

    public ImageResult createImageVariation(CreateImageVariationRequest request, java.io.File image) {
        RequestBody imageBody = RequestBody.create(MediaType.parse("image"), image);

        MultipartBody.Builder builder = new MultipartBody.Builder()
                .setType(MediaType.get("multipart/form-data"))
                .addFormDataPart("size", request.getSize())
                .addFormDataPart("response_format", request.getResponseFormat())
                .addFormDataPart("image", "image", imageBody);

        if (request.getN() != null) {
            builder.addFormDataPart("n", request.getN().toString());
        }

        return execute(api.createImageVariation(builder.build()));
    }

    public ModerationResult createModeration(ModerationRequest request) {
        return execute(api.createModeration(request));
    }

    /**
     * Calls the Open AI api, returns the response, and parses error messages if the request fails
     */
    public static <T> T execute(Single<T> apiCall) {
        try {
            return apiCall.blockingGet();
        } catch (HttpException e) {
            try {
                if (e.response() == null || e.response().errorBody() == null) {
                    throw e;
                }
                String errorBody = e.response().errorBody().string();

                OpenAiError error = mapper.readValue(errorBody, OpenAiError.class);
                throw new OpenAiHttpException(error, e, e.code());
            } catch (IOException ex) {
                // couldn't parse OpenAI error
                throw e;
            }
        }
    }

    /**
     * Calls the Open AI api and returns a Flowable of SSE for streaming
     * omitting the last message.
     *
     * @param apiCall The api call
     */
    public static Flowable<AzureSSE> stream(retrofit2.Call<ResponseBody> apiCall) {
        return stream(apiCall, false);
    }

    /**
     * Calls the Open AI api and returns a Flowable of SSE for streaming.
     *
     * @param apiCall  The api call
     * @param emitDone If true the last message ([DONE]) is emitted
     */
    public static Flowable<AzureSSE> stream(retrofit2.Call<ResponseBody> apiCall, boolean emitDone) {
        return Flowable.create(emitter -> apiCall.enqueue(new AzureResponseBodyCallback(emitter, emitDone)), BackpressureStrategy.BUFFER);
    }

    /**
     * Calls the Open AI api and returns a Flowable of type T for streaming
     * omitting the last message.
     *
     * @param apiCall The api call
     * @param cl      Class of type T to return
     */
    public static <T> Flowable<T> stream(Call<ResponseBody> apiCall, Class<T> cl) {
        return stream(apiCall).map(sse -> mapper.readValue(sse.getData(), cl));
    }

    /**
     * Shuts down the OkHttp ExecutorService.
     * The default behaviour of OkHttp's ExecutorService (ConnectionPool)
     * is to shut down after an idle timeout of 60s.
     * Call this method to shut down the ExecutorService immediately.
     */
    public void shutdownExecutor() {
        Objects.requireNonNull(this.executorService, "executorService must be set in order to shut down");
        this.executorService.shutdown();
    }


}
